use std::sync::Arc;
use std::collections::HashMap;
use tokio::sync::RwLock;
use serde::{Serialize, Deserialize};
use async_trait::async_trait;
use crate::workflow::errors::WorkflowError;

/// Represents the current state of a workflow node
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum NodeState {
    Pending,
    Running,
    Completed,
    Failed,
}

/// Represents the overall status of a workflow
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum WorkflowStatus {
    Created,
    Running,
    Completed,
    Failed,
    Cancelled,
}

/// Manages the state of workflow nodes
pub struct StateManager {
    backend: Arc<dyn StateBackend>,
    current_state: Arc<RwLock<HashMap<String, NodeState>>>,
}

impl StateManager {
    /// Creates a new StateManager with the specified backend
    pub fn new(backend: Arc<dyn StateBackend>) -> Self {
        Self {
            backend,
            current_state: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Updates the state of a node
    pub async fn update_node_state(&self, node_id: &str, state: NodeState) -> Result<(), WorkflowError> {
        let mut states = self.current_state.write().await;
        states.insert(node_id.to_string(), state.clone());
        self.backend.save_state(node_id, &state).await
    }

    /// Gets the current state of a node
    pub async fn get_node_state(&self, node_id: &str) -> Option<NodeState> {
        self.current_state.read().await.get(node_id).cloned()
    }
}

/// Trait for implementing different state storage backends
#[async_trait]
pub trait StateBackend: Send + Sync {
    async fn save_state(&self, node_id: &str, state: &NodeState) -> Result<(), WorkflowError>;
    async fn load_state(&self, node_id: &str) -> Result<Option<NodeState>, WorkflowError>;
} 